import argparse
import json
import os,utils
import tqdm
from datetime import datetime
import logging
import asyncio
import argparse
import json
import os
import asyncio
import aiohttp
import tqdm.asyncio as tqdm_async
from typing import Dict, Any
from tqdm import tqdm  # 使用标准tqdm而不是异步版本

# 禁用所有不必要的日志
logging.getLogger("openai").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)


global project_name 


# 并发数
semaphore = asyncio.Semaphore(10)  # 最大并发数

example = '''example 1:
    question:
    static int ossl_statem_server13_read_transition(SSL_CONNECTION *s, int mt)
{
    OSSL_STATEM *st = &s->statem;

    /*
     * Note: There is no case for TLS_ST_BEFORE because at that stage we have
     * not negotiated TLSv1.3 yet, so that case is handled by
     * ossl_statem_server_read_transition()
     */

    switch (st->hand_state) {
    default:
        break;

    case TLS_ST_EARLY_DATA:
        if (s->hello_retry_request == SSL_HRR_PENDING) {
            if (mt == SSL3_MT_CLIENT_HELLO) {
                st->hand_state = TLS_ST_SR_CLNT_HELLO;
                return 1;
            }

    case TLS_ST_CR_CERT_REQ:
        if (s->post_handshake_auth == SSL_PHA_REQUESTED) {
            if (do_compressed_cert(s))
                st->hand_state = TLS_ST_CW_COMP_CERT;
            else
                st->hand_state = TLS_ST_CW_CERT;
            return WRITE_TRAN_CONTINUE;
        
    Reasoning process:
    In the switch statement, the second case checks if TLS_ST_EARLY_DATA is one of the values in the given state list. Therefore, the initial state is TLS_ST_EARLY_DATA. Within this case, there are two nested if statements: if (s->hello_retry_request == SSL_HRR_PENDING) and if (mt == SSL3_MT_CLIENT_HELLO). Consequently, the condition is "s->hello_retry_request == SSL_HRR_PENDING AND mt == SSL3_MT_CLIENT_HELLO." Additionally, since st->hand_state = TLS_ST_SR_CLNT_HELLO, the final state is TLS_ST_SR_CLNT_HELLO.

    answer:

{
    "function": "ossl_statem_server13_read_transition",
    "transitions": [
        {
            "start_state": "TLS_ST_EARLY_DATA",
            "condition": "s->hello_retry_request == SSL_HRR_PENDING and mt == SSL3_MT_CLIENT_HELLO",
            "condition explanation": "server need client send Hello Retry Request __AND__ server read message is SSL3_MT_CLIENT_HELLO"
            "end_state": "TLS_ST_SR_CLNT_HELLO"
        },
        {
            "start_state": "TLS_ST_CR_CERT_REQ",
            "condition": "s->post_handshake_auth == SSL_PHA_REQUESTED and do_compressed_cert(s)",
            "condition explanation":"post handshake authentication was requested __AND__ compressed certificate was requested",
            "end_state": "TLS_ST_CW_COMP_CERT"
        },
        {
            "start_state": "TLS_ST_CR_CERT_REQ",
            "condition": "s->post_handshake_auth == SSL_PHA_REQUESTED and not do_compressed_cert(s)",
            "condition explanation":"post handshake authentication was requested __AND__ __NOT__ compressed certificate was requested",
            "end_state": "TLS_ST_CW_CERT"
        }
    ]
}


            '''


# get_fsm_prompt = lambda x,y: f'You are an experienced protocol expert. The protocol implementation is {x} Provide a function to extract state transitions from it, ensuring that all transitions are from one state to another, rather than a list of states. This point is crucial. The output format should be: {{"start_state":start state,"condition":condition of transition,"condition explanation":condition explanation in English,"end_state":end state}}.  The logical connective in condition explanation should be replaced with the following dictionary {{"and":"__AND__", "or":"__OR__","not":"__NOT__"}}  think step by step. {example}, \n\n question : {y}'  


get_fsm_prompt = lambda x,y,z,s: f'您是一位经验丰富的协议专家。协议实现是{x}。请提供一个函数来从中提取状态转换，确保所有转换都是从一种状态到另一种状态，而不是状态列表。这一点至关重要。输出格式应为：{{"start_state":起始状态,"condition":转换条件,"condition explanation":英文条件解释,"end_state":结束状态}}。条件解释中的逻辑连接词应替换为以下字典：{{"and":"__AND__", "or":"__OR__","not":"__NOT__"}}。请逐步思考。状态变量是：{z},状态变量取值列表是:{s} {example}, \n\n 问题：{y}'

def setup_logger(log_path):
    """Set up logging configuration"""
    # 创建 logger 对象（而不是直接用 basicConfig）
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)  # 全局最低级别设为 DEBUG
    
    # 清除默认的 handlers（避免重复日志）
    logger.handlers.clear()
    
    # 1. 文件处理器（记录 DEBUG 及以上级别）
    file_handler = logging.FileHandler(log_path, encoding="utf-8")
    file_handler.setLevel(logging.DEBUG)  # 文件记录所有 DEBUG+ 日志
    file_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(file_formatter)
    
    # 2. 控制台处理器（仅记录 INFO 及以上级别）
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)  # 终端只打印 INFO+
    console_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 添加处理器
    logger.addHandler(file_handler)
    # logger.addHandler(console_handler)
    
    return logger
            


def process_function(group_name: str, func_name: str, func_data: Dict, 
                          protocol_impl: str, logger,state_vars,state_var_values):
    """处理单个函数的异步任务"""


    try:
        prompt = get_fsm_prompt(protocol_impl, func_data, state_vars, state_var_values)
        logger.info(f"Sending to LLM: {func_name}")
        
        answer = utils.get_ans_json(prompt)
        
        logger.debug(f"Response for {func_name}:\n{json.dumps(answer, indent=2)}")
        return (group_name, func_name, answer)
        
    except Exception as e:
        logger.error(f"Error in {func_name}: {str(e)}")
        return (group_name, func_name, {"error": str(e)})


def get_function(config, func_name):
    """
    从配置指定的项目中获取特定函数的源代码
    
    Args:
        config: 包含项目配置的字典
        func_name: 要提取的函数名称
    
    Returns:
        str: 函数的源代码内容，如果找不到则返回None
    """
    try:
        # 获取项目目录和函数信息
        project_dir = config.get("project_dir")
        state_functions = config.get("state_functions", {})
        
        # 在所有状态函数中查找目标函数
        target_func_info = None
        for state_name, functions in state_functions.items():
            if func_name in functions:
                target_func_info = functions[func_name]
                break
        
        if not target_func_info:
            print(f"函数 {func_name} 未在配置中找到")
            return None
        
        # 构建完整文件路径
        relative_path = target_func_info["path"]
        full_path = os.path.join(project_dir, relative_path)
        
        # 检查文件是否存在
        if not os.path.exists(full_path):
            print(f"文件不存在: {full_path}")
            return None
        
        # 读取指定行号范围的内容
        start_line = target_func_info["start"]
        end_line = target_func_info["end"]
        
        return extract_code_snippet(full_path, start_line, end_line)
        
    except Exception as e:
        print(f"获取函数时出错: {str(e)}")
        return None

def concise_function(state_var, function_code):
    """
    检查函数代码中是否有对state_var的赋值操作，
    并返回包含这些赋值的代码片段（智能扩展切片范围，最多100行）。
    
    参数:
        state_var: 要检查的变量名（如'ssl->options.serverState'）
        function_code: 函数的源代码字符串
        
    返回:
        生成器，每次生成一个包含赋值操作的代码片段
    """
    lines = function_code.split('\n')
    var_parts = state_var.replace('->', '.').split('.')
    var_name = var_parts[-1]  # 提取变量名最后部分
    
    # 找到所有赋值操作的行号
    assignment_lines = []
    for i, line in enumerate(lines):
        stripped = line.strip()
        # 跳过注释行
        if stripped.startswith('#') or stripped.startswith('//'):
            continue
            
        # 检查简单赋值
        if '=' in line and any(
            part in line.split('=')[0] 
            for part in var_parts
        ):
            assignment_lines.append(i)
        
        # 检查复合赋值
        for op in ['+=', '-=']:
            if op in line and any(
                part in line.split(op)[0] 
                for part in var_parts
            ):
                assignment_lines.append(i)
                break
    
    # 合并相邻的赋值区域
    merged_ranges = []
    if assignment_lines:
        current_start = assignment_lines[0]
        current_end = assignment_lines[0]
        
        for line in assignment_lines[1:]:
            if line <= current_end + 50:  # 如果下一个赋值在50行内，扩展范围
                current_end = line
            else:
                merged_ranges.append((current_start, current_end))
                current_start = current_end = line
        merged_ranges.append((current_start, current_end))
    
    # 生成代码片段
    for start, end in merged_ranges:
        # 计算扩展范围，确保不超过100行
        context_size = min(25, 100 // (end - start + 1)) if start != end else 25
        expanded_start = max(0, start - context_size)
        expanded_end = min(len(lines), end + context_size + 1)
        
        # 确保总行数不超过100
        while expanded_end - expanded_start > 100:
            if expanded_end - start > end - expanded_start:
                expanded_end -= 1
            else:
                expanded_start += 1
        
        yield '\n'.join(lines[expanded_start:expanded_end])



def extract_code_snippet(file_path, start_line, end_line):
    """
    从文件中提取指定行号范围的代码
    
    Args:
        file_path: 文件路径
        start_line: 起始行号
        end_line: 结束行号
    
    Returns:
        str: 提取的代码内容
    """
    try:
        extracted_lines = []
        
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as file:
            for current_line_num, line in enumerate(file, 1):
                if current_line_num >= start_line and current_line_num <= end_line:
                    extracted_lines.append(line)
                elif current_line_num > end_line:
                    break
        
        return ''.join(extracted_lines)
        
    except FileNotFoundError:
        print(f"文件未找到: {file_path}")
        return None
    except Exception as e:
        print(f"读取文件时出错: {str(e)}")
        return None






def main(config):
    """同步主函数"""
    save_dir = os.path.join(os.path.abspath(config['save_parent_dir']), 
                          f"{project_name}_{config['version']}")
    os.makedirs(save_dir, exist_ok=True)
    
    # 初始化日志
    log_path = os.path.join(save_dir, 'get_fsm_from_llm.log')
    logger = setup_logger(log_path)
    logger.info(f"Starting processing for {project_name}")
    
    # 准备任务
    results = []
    for group_name, func_names in tqdm(config['state_functions'].items(), desc="Processing groups"):
        state_var = None
        state_enum = None
        for func_name in func_names:
            # 跳过前两个特殊数据
            if func_name == 'state_var':
                state_var = config['state_functions'][group_name][func_name]
                continue
            if func_name == 'state_enum':
                state_enum = config['state_functions'][group_name][func_name]
                continue
            
            for snippet in concise_function(state_var, get_function(config, func_name)):
                # 直接调用同步版本的process_function
                result = process_function(
                    group_name, func_name, 
                    snippet, config['project'],
                    logger, state_var, state_enum
                )
                results.append(result)
    
    # 整理结果
    llm_fsm_response = {}
    for group_name, func_name, answer in results:
        if group_name not in llm_fsm_response:
            llm_fsm_response[group_name] = {}
        llm_fsm_response[group_name][func_name] = answer
    
    # 保存结果
    save_path = os.path.join(save_dir, "get_fsm_from_llm.json")
    with open(save_path, 'w', encoding='utf-8') as f:
        json.dump(llm_fsm_response, f, indent=2)
    
    logger.info(f"Saved results to {save_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Sync code analysis')
    parser.add_argument('config', help='Path to config JSON')
    config = parser.parse_args().config

    # config = r".\config\test.json"

    with open(config, encoding='utf-8') as f:
        config = json.load(f)

    project_name = config['project']
    
    # 直接调用同步主函数
    main(config)